import os
from datetime import datetime, timedelta
import asyncio
import pandas as pd
from tortoise import Tortoise
from models import Stock, StockPrice
import time
import akshare as ak 


def get_default_year():
    return datetime.now().year

async def init_db(year=get_default_year()):
    """初始化数据库连接并创建表"""
    try:
        # 使用绝对路径
        BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        DB_PATH = os.path.join(BASE_DIR, f"{year}_db.sqlite3")
        
        # 初始化数据库连接
        await Tortoise.init(
            db_url=f"sqlite://{DB_PATH}",
            modules={"models": ["models"]}
        )
        
        # 生成数据库表
        await Tortoise.generate_schemas(safe=True)
        print("Database initialized successfully")
        return True
    except Exception as e:
        print(f"Database initialization failed: {e}")
        return False

def get_and_process_data(symbol, start_date, end_date, adjust):
    """获取并处理数据，包括重命名列等操作"""
    df = ak.stock_zh_a_hist(
        symbol=symbol,
        period="daily",
        start_date=start_date.strftime("%Y%m%d"),
        end_date=end_date.strftime("%Y%m%d"),
        adjust=adjust
    )

    if df is None or df.empty:
        return pd.DataFrame()

    # 重命名列以匹配技术分析代理的需求
    df = df.rename(columns={
        "日期": "date",
        "开盘": "open",
        "最高": "high",
        "最低": "low",
        "收盘": "close",
        "成交量": "volume",
        "成交额": "amount",
        "振幅": "amplitude",
        "涨跌幅": "pct_change",
        "涨跌额": "change_amount",
        "换手率": "turnover"
    })

    # 确保日期列为datetime类型
    df["date"] = pd.to_datetime(df["date"])
    return df


async def get_and_process_data_fromdb(symbol, start_date, end_date, adjust="qfq"):
    """
    从本地数据库获取股票历史数据
    :param symbol: 股票代码
    :param start_date: 开始日期
    :param end_date: 结束日期
    :param adjust: 复权类型，保持参数兼容性
    :return: 处理后的股票数据DataFrame
    """
    try:
        await init_db()
        
        # 查询数据
        query = StockPrice.filter(symbol=symbol)
        if start_date:
            query = query.filter(date__gte=start_date)
        if end_date:
            query = query.filter(date__lte=end_date)
        records = await query
        
        # 转换为DataFrame
        df = pd.DataFrame([{
            "symbol": r.symbol,
            "date": r.date,
            "open": float(r.open_price),  # 转换为float
            "high": float(r.high_price),   # 转换为float
            "low": float(r.low_price),     # 转换为float
            "close": float(r.close_price), # 转换为float
            "volume": float(r.volume),     # 转换为float
            "amount": float(r.amount),     # 转换为float
            "amplitude": float(r.amplitude), # 转换为float
            "pct_change": float(r.pct_change), # 转换为float
            "change_amount": float(r.change_amount), # 转换为float
            "turnover": float(r.turnover)  # 转换为float
        } for r in records])
        
        # 按日期排序
        df = df.sort_values('date')
        
        # 关闭数据库连接
        await Tortoise.close_connections()
        
        return df
        
    except Exception as e:
        print(f"Error getting data from database for {symbol}: {e}")
        await Tortoise.close_connections()
        return pd.DataFrame()


def get_and_process_data_fromdb_sync(symbol, start_date, end_date, adjust="qfq"):
    """
    同步方式从本地数据库获取股票历史数据
    :param symbol: 股票代码
    :param start_date: 开始日期
    :param end_date: 结束日期
    :param adjust: 复权类型，保持参数兼容性
    :return: 处理后的股票数据DataFrame
    """
    try:
        # 获取或创建事件循环
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        # 运行异步函数
        return loop.run_until_complete(
            get_and_process_data_fromdb(symbol, start_date, end_date, adjust)
        )
    except Exception as e:
        print(f"Error in sync version of get_and_process_data_fromdb: {e}")
        return pd.DataFrame()


async def save_price_data(price_data):
    price_records = []
    for _, row in price_data.iterrows():
        price_records.append({
            "symbol": row["symbol"],
            "date": row["date"],
            "open_price": row["open"],
            "high_price": row["high"],
            "low_price": row["low"],
            "close_price": row["close"],
            "volume": row["volume"],
            "amount": row["amount"],
            "amplitude": row["amplitude"],
            "pct_change": row["pct_change"],
            "change_amount": row["change_amount"],
            "turnover": row["turnover"],
        })
    if len(price_records) == 0:
        return
    # 批量插入
    update_fields = [
        "open_price", "close_price", "high_price", "low_price",
        "volume", "amount", "amplitude", "pct_change", "change_amount", "turnover"]
    await StockPrice.bulk_create(
        [StockPrice(**record) for record in price_records],
        ignore_conflicts=True, on_conflict=["symbol", "date"])

async def migrate_data(stock_symbols, start_date=datetime.now()-timedelta(days=365), end_date=datetime.now()):
    """迁移现有数据到数据库"""
    for symbol in stock_symbols:
        
        try:
            price_data = get_and_process_data(symbol, start_date, end_date, "qfq")
            if not price_data.empty:
                # 准备批量插入的数据
                print(f"Migrated price data for {symbol}, count: {len(price_records)}")
                time.sleep(5)
                price_data["symbol"] = [symbol for _ in range(len(price_data))]
                await save_price_data(price_data)
        except Exception as e:
            print(f"Error migrating price data for {symbol}: {e}")
    
    await Tortoise.close_connections()


async def check_tables():
    conn = Tortoise.get_connection("default")
    tables = await conn.execute_query("SELECT name FROM sqlite_master WHERE type='table'")
    print("Tables in database:", [table[0] for table in tables[1]])


async def check_symbols():
    symbols = await StockPrice.all().values_list('symbol', flat=True)
    return  set(symbols)


async def get_all_symbols():
    r = ak.stock_zh_a_spot_em()

    df = r.rename(columns={
        "代码": "symbol",
        "今开": "open",
        "最高": "high",
        "最低": "low",
        "最新价": "close",
        "成交量": "volume",
        "成交额": "amount",
        "振幅": "amplitude",
        "涨跌幅": "pct_change",
        "涨跌额": "change_amount",
        "换手率": "turnover"
    })
    records = []
    index_df = ak.stock_zh_a_daily(symbol="sh000001", start_date=(datetime.now()-timedelta(days=30)).strftime("%Y%m%d"))
    latest_date = index_df.iloc[-1]['date']
    # 获取最新的交易日期
    df['date'] = [latest_date for _ in range(len(df))]
    need_realtime_data = []
    for _, row in df.iterrows():
        if await StockPrice.filter(symbol=row['symbol'], date=row['date']).exists():
            continue
        need_realtime_data.append(dict(row))
    need_realtime_df = pd.DataFrame(need_realtime_data)
    print(need_realtime_df)
    await save_price_data(need_realtime_df)

    exists_symbols = await Stock.filter().all().values_list('symbol', flat=True)

    for idx, row in r.iterrows():
        symbol = row['代码']
        name = row['名称']
        records.append(symbol)
        if symbol in exists_symbols:
            continue

        stock, created = await Stock.get_or_create(
                    symbol=symbol,
                    defaults={"name": name}
                )
    return set(records)

async def get_symbol_data(symbol, start=datetime(2025,1,1), end=datetime.now()):
    year = start.year
    await init_db(year)
    query = StockPrice.filter(symbol=symbol)
    if start:
        query = query.filter(date__gte=start)
    if end:
        query = query.filter(date__lte=end)
    records = await query
    await Tortoise.close_connections()
    
    df = pd.DataFrame([{
        "代码": r.symbol,
        "日期": r.date,
        "今开": r.open_price,
        "最高": r.high_price,
        "最低": r.low_price,
        "收盘": r.close_price,
        "成交量": r.volume,
        "成交额": r.amount,
        "振幅": r.amplitude,
        "涨跌幅": r.pct_change,
        "涨跌额": r.change_amount,
        "换手率": r.turnover
    } for r in records])
    return df

async def main(year):

    # 初始化数据库
    ret = await init_db(year)
    if not ret:
        print("Failed to initialize database, exiting...")
        return
    all_symbols = await get_all_symbols()
    exists_symbols = await check_symbols()
    needed = all_symbols - exists_symbols
    print(len(needed))
    start = datetime(year, 1, 1)
    end = datetime(year+1, 1, 1)
    #await migrate_data(list(needed), start, end)
    await Tortoise.close_connections()

if __name__ == "__main__":
    # TODO 1.实现按照年份的分库分表；2.实现数据的插入逻辑
    asyncio.run(main(2025))
    
    #r = asyncio.run(get_symbol_data('300400'))
    #print(r)
